(*
   Copyright (C) 2024 International Digital Economy Academy.
   This program is licensed under the MoonBit Public Source
   License as published by the International Digital Economy Academy,
   either version 1 of the License, or (at your option) any later
   version. This program is distributed in the hope that it will be
   useful, but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the MoonBit
   Public Source License for more details. You should have received a
   copy of the MoonBit Public Source License along with this program. If
   not, see
   <https://www.moonbitlang.com/licenses/moonbit-public-source-license-v1>.
*)

(* 内存池模块，用于优化编译器的内存使用 *)

(* 内存块大小 *)
let default_block_size = 4096 (* 4KB *)

(* 内存池类型 *)
type 'a t = {
  mutable blocks: 'a array list;  (* 内存块列表 *)
  mutable current_block: 'a array;  (* 当前使用的内存块 *)
  mutable current_index: int;  (* 当前内存块的使用索引 *)
  block_size: int;  (* 内存块大小 *)
  create_item: unit -> 'a;  (* 创建新项的函数 *)
}

(* 创建新的内存池 *)
let create ?(block_size = default_block_size) ~create_item () =
  let initial_block = Array.make block_size (create_item ()) in
  {
    blocks = [initial_block];
    current_block = initial_block;
    current_index = 0;
    block_size;
    create_item;
  }

(* 从内存池中获取一个项 *)
let get pool =
  if pool.current_index >= pool.block_size then begin
    (* 当前块已满，创建新块或重用已有块 *)
    let next_block = 
      try List.hd (List.tl pool.blocks) with
      | _ -> 
          (* 没有可重用的块，创建新块 *)
          let new_block = Array.make pool.block_size (pool.create_item ()) in
          pool.blocks <- new_block :: pool.blocks;
          new_block
    in
    pool.current_block <- next_block;
    pool.current_index <- 0;
  end;
  
  let item = pool.current_block.(pool.current_index) in
  pool.current_index <- pool.current_index + 1;
  item

(* 重置内存池，但不释放内存 *)
let reset pool =
  pool.current_block <- List.hd pool.blocks;
  pool.current_index <- 0

(* 清空内存池，释放所有内存 *)
let clear pool =
  if List.length pool.blocks > 0 then begin
    pool.current_block <- List.hd pool.blocks;
    pool.blocks <- [pool.current_block];
    pool.current_index <- 0;
    (* 重新初始化当前块 *)
    for i = 0 to pool.block_size - 1 do
      pool.current_block.(i) <- pool.create_item ()
    done;
  end

(* 获取内存池的统计信息 *)
let stats pool =
  let total_blocks = List.length pool.blocks in
  let total_capacity = total_blocks * pool.block_size in
  let used_items = ((total_blocks - 1) * pool.block_size) + pool.current_index in
  (total_blocks, total_capacity, used_items)

(* 内存池迭代器 *)
let iter f pool =
  let rec iter_blocks = function
    | [] -> ()
    | block :: rest ->
        if block == pool.current_block then
          (* 当前块只迭代已使用的部分 *)
          for i = 0 to pool.current_index - 1 do
            f block.(i)
          done
        else
          (* 其他块迭代全部 *)
          Array.iter f block;
        iter_blocks rest
  in
  iter_blocks pool.blocks